# -*- coding: utf-8 -*-
"""
Generators which take a bundle and return an interator over the increasing
closure of the provided bundle according to a given preorder.

the bundle c is assumed to have four elements and 
c[0] represents state 1 - time 1
c[1] represents state 1 - time 2
c[2] represents state 2 - time 1
c[3] represents state 2 - time 2
"""

import numpy as np

def lottery_equivalence(c):
    yield c
    yield np.array( [ c[2],c[3],c[0],c[1] ] )

def time_neutrality(c):
    yield c
    yield np.array( [ c[1], c[0], c[3], c[2] ] )

def impatience(c):
    yield c
    if c[1] > c[0]:
        yield np.array([ c[1],c[0],c[2],c[3] ])
    if c[3] > c[2]:
        yield np.array( [ c[0],c[1],c[3],c[2] ] )
    if (c[1] > c[0]) & (c[3] > c[2]):
        yield np.array( [ c[1],c[0],c[3],c[2] ] )

def patience(c):
    yield c
    if c[0] > c[1]:
        yield np.array([c[1],c[0],c[2],c[3]])
    if c[2] > c[3]:
        yield np.array([c[0],c[1],c[3],c[2]])
    if ( c[0] > c[1] ) & (c[2] > c[3]):
        yield np.array( [c[1],c[0],c[3],c[2]] )

def correlation_neutrality(c):
    yield c
    yield np.array( [c[2],c[1],c[0],c[3]] )
    yield np.array( [c[0],c[3],c[2],c[1]] )

def correlation_aversion(c):
    yield c
    if ( (c[0] >= c[2]) & (c[1] >= c[3]) ) | ( (c[0] <= c[2]) & (c[1] <= c[3]) ) :
        yield np.array([ c[2],c[1],c[0],c[3] ])
        yield np.array([ c[0],c[3],c[2],c[1] ])
        
def correlation_seeking(c):
    yield c
    if ( (c[0] >= c[2]) & (c[1] <= c[3]) ) | ( (c[0] <= c[2]) & (c[1] >= c[3]) ) :
        yield np.array([ c[2],c[1],c[0],c[3] ])
        yield np.array([ c[0],c[3],c[2],c[1] ])

def stochastic_impatience(c):
    yield c
    diag = max( c[0],c[3] )
    if ( c[1] >= diag ) & ( c[2] >= diag ) & ( c[2] < c[1] ):
        yield np.array( [diag, c[2], c[1], diag ] )
    diag = max( c[1],c[2] )
    if ( c[0] >= diag ) & ( c[3] >= diag ) & ( c[0] < c[3] ):
        yield np.array( [c[3],diag,diag,c[0]] )

def stochastic_patience(c):
    yield c
    diag = max( c[0],c[3] )
    if ( c[1] >= diag ) & ( c[2] >= diag ) & ( c[2] > c[1] ):
        yield np.array( [diag, c[2], c[1], diag ] )
    diag = max( c[1],c[2] )
    if ( c[0] >= diag ) & ( c[3] >= diag ) & ( c[0] > c[3] ):
        yield np.array( [c[3],diag,diag,c[0]] )

def impatience2(c):
    '''Two good version of impatience.'''
    yield c
    if c[1] > c[0]:
        yield np.array( [c[1],c[0]] )
    
def patience2(c):
    '''Two good version of patience.'''
    yield c
    if c[0] > c[1]:
        yield np.array( [c[1],c[0]] )
    
    

